@file:Suppress("unused")

package cn.qfys521.mc_websocket.client

import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.serialization.KSerializer
import kotlinx.serialization.json.*
import org.java_websocket.client.WebSocketClient
import org.java_websocket.handshake.ServerHandshake
import java.net.URI
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicLong
import kotlinx.serialization.descriptors.SerialDescriptor
import kotlinx.serialization.descriptors.buildClassSerialDescriptor
import kotlinx.serialization.encoding.Decoder
import kotlinx.serialization.encoding.Encoder
import kotlin.random.Random

/**
 * 基于 WebSocket 的 JSON-RPC 2.0 客户端，适配 Minecraft 服务端管理协议。
 *
 * 特性：
 * - 自动重连（指数退避+抖动，可配置）
 * - 请求发送失败自动重试
 * - 通知分发（方法名完整，如："notification:players/joined"）
 *
 * 线程模型：
 * - 内部使用协程调度（IO），通知回调在内部协程中触发。
 * - [call] 与 [callNoParams] 为挂起函数，等待响应后返回或抛出 [ManagementException]。
 *
 * 错误：见 [ManagementException] 层次结构。
 */
class ManagementWsClient(
    serverUri: URI,
    headers: Map<String, String> = emptyMap(),
    private val config: ReconnectConfig = ReconnectConfig()
) : WebSocketClient(serverUri, headers) {

    /**
     * 自动重连与重试配置。
     *
     * @property autoReconnect 是否自动重连（连接关闭/错误时）。
     * @property reconnectBaseDelayMs 重连初始等待毫秒数（指数退避基数）。
     * @property reconnectMaxDelayMs 重连最大等待毫秒数。
     * @property reconnectJitterMs 重连等待抖动上限毫秒数（0..jitter）。
     * @property connectWaitTimeoutMs [ensureOpen] 默认等待超时时间毫秒数。
     * @property callSendRetry [call] 发送失败或未连接时的最大重试次数。
     * @property callRetryBaseDelayMs 调用重试的初始退避毫秒数。
     * @property callRetryMaxDelayMs 调用重试的最大退避毫秒数。
     */
    data class ReconnectConfig(
        val autoReconnect: Boolean = true,
        val reconnectBaseDelayMs: Long = 500,
        val reconnectMaxDelayMs: Long = 10_000,
        val reconnectJitterMs: Long = 250,
        val connectWaitTimeoutMs: Long = 5_000,
        val callSendRetry: Int = 3,
        val callRetryBaseDelayMs: Long = 200,
        val callRetryMaxDelayMs: Long = 1_000,
    )

    private val scope = CoroutineScope(Dispatchers.IO + Job())
    private val json = Json { ignoreUnknownKeys = true; encodeDefaults = false; explicitNulls = false }

    private val nextId = AtomicLong(1)
    private val pending = ConcurrentHashMap<Long, CompletableDeferred<String>>()
    private val notificationListeners = ConcurrentHashMap<String, MutableList<(JsonElement?) -> Unit>>()

    private val openSignalMutex = Mutex()
    private var openSignal: CompletableDeferred<Unit> = CompletableDeferred()

    private var reconnectJob: Job? = null

    override fun onOpen(handshakedata: ServerHandshake?) {
        scope.launch {
            openSignalMutex.withLock {
                if (!openSignal.isCompleted) openSignal.complete(Unit)
            }
        }
    }

    override fun onClose(code: Int, reason: String?, remote: Boolean) {
        
        val ex = ManagementException.ConnectionClosed(code, reason, remote)
        pending.values.forEach { it.completeExceptionally(ex) }
        pending.clear()
        scope.launch {
            openSignalMutex.withLock {
                openSignal = CompletableDeferred()
            }
        }
        if (config.autoReconnect) scheduleReconnect()
    }

    override fun onError(ex: Exception?) {
        
        if (config.autoReconnect) scheduleReconnect()
    }

    override fun onMessage(message: String) {
        
        val element = runCatching { json.parseToJsonElement(message) }.getOrNull() as? JsonObject ?: return

        val id = element["id"]
        if (id != null && id !is JsonNull) {
            val idLong = id.jsonPrimitive.long
            pending.remove(idLong)?.complete(message)
            return
        }

        val method = element["method"]?.jsonPrimitive?.contentOrNull
        if (!method.isNullOrEmpty()) {
            val params = element["params"]
            
            notificationListeners[method]?.forEach { listener ->
                scope.launch { listener(params) }
            }
        }
    }

    private fun scheduleReconnect() {
        if (isOpen) return
        if (reconnectJob?.isActive == true) return
        reconnectJob = scope.launch {
            var delayMs = config.reconnectBaseDelayMs
            while (!isOpen) {
                try {
                    
                    reconnectBlocking()
                } catch (_: Throwable) {
                    
                }
                if (isOpen) break
                val jitter = Random.nextLong(0, config.reconnectJitterMs + 1)
                delay(delayMs + jitter)
                delayMs = (delayMs * 2).coerceAtMost(config.reconnectMaxDelayMs)
            }
        }
    }

    /**
     * 确保连接处于已打开状态；若未打开则等待，必要时触发自动重连（仅在 isClosed==true 时）。
     *
     * @param timeoutMs 最长等待毫秒数，默认使用 [ReconnectConfig.connectWaitTimeoutMs]。
     * @throws ManagementException.Timeout 超时仍未打开。
     */
    suspend fun ensureOpen(timeoutMs: Long = config.connectWaitTimeoutMs) {
        if (isOpen) return
        
        if (config.autoReconnect && isClosed) scheduleReconnect()
        
        val wait = openSignalMutex.withLock { openSignal }
        try {
            withTimeoutK(timeoutMs) { wait.await() }
        } catch (t: Throwable) {
            throw ManagementException.Timeout("Timed out waiting for connection open ($timeoutMs ms)")
        }
    }

    /**
     * 订阅指定通知方法，自动将 params 单元素数组解包并反序列化为类型 [P] 后回调。
     *
     * @param method 通知方法名（形如："notification:server/status"）。
     * @param paramsDeserializer 负载反序列化器。
     * @param callback 收到通知时的回调（可能为 null，表示无/空负载）。
     */
    fun <P> onNotification(
        method: String,
        paramsDeserializer: KSerializer<P>,
        callback: (P?) -> Unit
    ) {
        notificationListeners.computeIfAbsent(method) { mutableListOf() }
            .add { paramsElement ->
                val payload: P? = paramsElement?.let { unwrapParams(it, paramsDeserializer) }
                callback(payload)
            }
    }

    /**
     * 订阅指定通知方法，提供原始 JSON 负载（未解包/未反序列化）。
     *
     * @param method 通知方法名。
     * @param callback 原始 JSON 负载（可能为 null）。
     */
    fun onRawNotification(method: String, callback: (JsonElement?) -> Unit) {
        notificationListeners.computeIfAbsent(method) { mutableListOf() }.add(callback)
    }

    /**
     * 发起带参数的 JSON-RPC 调用（params 将按协议封装为单元素数组）。
     *
     * @param method 方法名（如："minecraft:server/status"）。
     * @param params 参数对象；序列化后作为 params[0] 发送。
     * @param paramsSerializer 参数序列化器。
     * @param resultDeserializer 结果反序列化器。
     * @return 反序列化后的结果对象。
     *
     * 错误：
     * - [ManagementException.ConnectionClosed] 连接关闭/未连接。
     * - [ManagementException.Timeout] 等待打开或响应超时（间接）。
     * - [ManagementException.Rpc] 服务端返回错误。
     * - [ManagementException.ParseError] 响应解析失败。
     * - [ManagementException.Transport] 其他底层错误。
     */
    suspend fun <P, R> call(
        method: String,
        params: P?,
        paramsSerializer: KSerializer<P>,
        resultDeserializer: KSerializer<R>
    ): R {
        val id = nextId.getAndIncrement()
        val reqObj = buildRequest(method, id, params, paramsSerializer)
        val reqText = composeRequestText(reqObj)

        var attempt = 0
        var lastError: Throwable? = null
        while (attempt < config.callSendRetry) {
            attempt++
            try {
                ensureOpen()
                val waiter = CompletableDeferred<String>()
                pending[id] = waiter
                send(reqText)
                val respText = waiter.await()
                return decodeResultFromResponse(resultDeserializer, respText)
            } catch (t: Throwable) {
                pending.remove(id)
                lastError = t
                if (attempt >= config.callSendRetry) break
                
                val backoff = backoffForAttempt(attempt, config.callRetryBaseDelayMs, config.callRetryMaxDelayMs)
                delay(backoff)
            }
        }
        
        throw when (lastError) {
            is ManagementException -> lastError
            is java.nio.channels.NotYetConnectedException -> ManagementException.ConnectionClosed(null, "Not yet connected", null)
            else -> ManagementException.Transport(lastError ?: IllegalStateException("Unknown error"))
        }
    }

    /**
     * 发起无参数 JSON-RPC 调用（省略 params 字段）。
     *
     * @param method 方法名。
     * @param resultDeserializer 结果反序列化器。
     * @return 反序列化后的结果对象。
     * @see call
     */
    suspend fun <R> callNoParams(
        method: String,
        resultDeserializer: KSerializer<R>
    ): R = call<Unit, R>(method, null, NoParamsSerializer, resultDeserializer)

    internal fun composeRequestText(obj: JsonObject): String =
        json.encodeToString(JsonObject.serializer(), obj)

    internal fun <R> decodeResultFromResponse(resultDeserializer: KSerializer<R>, respText: String): R {
        val resp = runCatching { json.decodeFromString(JsonObject.serializer(), respText) }
            .getOrElse { throw ManagementException.ParseError("Invalid JSON response", it) }
        val error = resp["error"]
        if (error != null && error !is JsonNull) {
            val errObj = error.jsonObject
            val code = errObj["code"]?.jsonPrimitive?.intOrNull
            val message = errObj["message"]?.jsonPrimitive?.contentOrNull
            val data = errObj["data"]
            throw ManagementException.Rpc(code, message ?: "RPC error", data)
        }
        val resultEl = resp["result"] ?: throw ManagementException.ParseError("Missing result field")
        return runCatching { json.decodeFromJsonElement(resultDeserializer, resultEl) }
            .getOrElse { throw ManagementException.ParseError("Result decode failed", it) }
    }

    internal fun <P> buildRequest(
        method: String,
        id: Long,
        params: P?,
        paramsSerializer: KSerializer<P>
    ): JsonObject {
        val obj = buildMap<String, JsonElement> {
            put("jsonrpc", JsonPrimitive("2.0"))
            put("method", JsonPrimitive(method))
            put("id", JsonPrimitive(id))
            if (params != null) {
                put("params", wrapParams(json.encodeToJsonElement(paramsSerializer, params)))
            }
        }
        return JsonObject(obj)
    }

    private fun wrapParams(single: JsonElement): JsonArray = buildJsonArray { add(single) }

    private fun <P> unwrapParams(paramsEl: JsonElement, deserializer: KSerializer<P>): P? {
        
        return when (paramsEl) {
            is JsonArray -> paramsEl.firstOrNull()?.let { json.decodeFromJsonElement(deserializer, it) }
            is JsonObject -> json.decodeFromJsonElement(deserializer, paramsEl)
            else -> null
        }
    }

    private object NoParamsSerializer : KSerializer<Unit> {
        override val descriptor: SerialDescriptor = buildClassSerialDescriptor("kotlin.Unit")
        override fun serialize(encoder: Encoder, value: Unit) { /* no-op */ }
        override fun deserialize(decoder: Decoder): Unit = Unit
    }

    private suspend fun withTimeoutK(timeoutMs: Long, block: suspend () -> Unit) {
        
        val gate = CompletableDeferred<Unit>()
        val job = scope.launch { block(); gate.complete(Unit) }
        val result = scope.launch {
            delay(timeoutMs)
            if (!gate.isCompleted) gate.completeExceptionally(ManagementException.Timeout("Timeout $timeoutMs ms"))
        }
        try {
            gate.await()
        } finally {
            result.cancel()
            job.cancel()
        }
    }

    private fun backoffForAttempt(attempt: Int, base: Long, max: Long): Long {
        val exp = (base * (1L shl (attempt - 1))).coerceAtMost(max)
        val jitter = Random.nextLong(0, (base / 2).coerceAtLeast(1))
        return exp + jitter
    }
}
